import { houseScaleData } from "./dataset";
import { normalise, denormalise } from "./utils/TensorUtils";
import { createModel, trainModel } from "./models";
import { normalisedData } from "./types/normalisedData";
import { point } from "./types/point";
import { predict } from "./testModel/predicted";
import * as tf from "@tensorflow/tfjs-node";
import type {
  History,
  Sequential,
  Tensor,
  Tensor2D,
} from "@tensorflow/tfjs-node";

// 线性回归的主函数，是异步的
async function run(): Promise<string> {
  const points: point[] = (await houseScaleData.toArray()).map(
    (record: any) => {
      let sqft_living: number = record.sqft_living;
      let price: number = record.price;
      return Object({
        x: sqft_living,
        y: price,
      });
    }
  );

  // 这里的意义是如果数据集不能被切分所以会报错
  if (points.length % 2 !== 0) {
    points.pop();
  }

  // 这里是打乱数据集
  tf.util.shuffle(points);

  const featureValues: number[] = points.map((p: point) => p.x);
  const featureTensor: Tensor2D = tf.tensor2d(featureValues, [
    featureValues.length,
    1,
  ]);

  const labelValues: number[] = points.map((p: point) => p.y);
  const labelTensor: Tensor2D = tf.tensor2d(labelValues, [
    labelValues.length,
    1,
  ]);

  const normalisedFeature: normalisedData = normalise(featureTensor);
  const normalisedLabel: normalisedData = normalise(labelTensor);

  // denormalise(
  //   normalisedFeature.tensor,
  //   normalisedFeature.min,
  //   normalisedFeature.max
  // ).print();

  // 这里切分数据集
  const [trainNormalisedFeatureTensor, testNormalisedFeatureTensor]: Tensor[] =
    tf.split(normalisedFeature.tensor, 2, 0);

  const [trainNormalisedLabelTensor, testNormalisedLabelTensor]: Tensor[] =
    tf.split(normalisedLabel.tensor, 2, 0);

  // trainNormalisedFeatureTensor.print(true);

  const model: Sequential = createModel();
  model.summary();
  const result: History = await trainModel(
    model,
    trainNormalisedFeatureTensor,
    trainNormalisedLabelTensor
  );

  const trainingLoss: number | Tensor<tf.Rank> | undefined =
    result.history.loss.pop();
  console.log("training loss: " + trainingLoss);

  const validationLoss: number | Tensor<tf.Rank> | undefined =
    result.history.val_loss.pop();
  console.log("validating loss: " + validationLoss);

  const lossTensor: tf.Scalar | tf.Scalar[] = model.evaluate(
    testNormalisedFeatureTensor,
    testNormalisedLabelTensor
  );
  const loss: Uint8Array | Int32Array | Float32Array = (
    lossTensor as tf.Scalar
  ).dataSync();
  console.log("testing loss: " + loss);

  let outNumber: number = await predict(
    normalisedFeature,
    normalisedLabel,
    model
  );

  // 打印出预测的值
  console.log(outNumber);
  return "TFJS yyds";
}

export { run };
